package com.mydemo.controller;

import com.mydemo.vo.ProgressResponse;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.concurrent.atomic.AtomicLong;

@RestController
public class ChatClientV2Controller {

    @Resource
    private ChatClient chatClient;

    @GetMapping("/chat/client/v2")
    public Flux<String> hello2(@RequestParam(name="msg",defaultValue = "祖国万岁" )String msg){
        //普通流式相应
        Flux<String> result = chatClient.prompt(msg).stream().content();

        return result;
    }

    //1.添加计时功能
    @GetMapping("/chat/client/v2.1")
    public Flux<String> hello21(@RequestParam(name="msg",defaultValue = "祖国万岁" )String msg){

        long startTime = System.currentTimeMillis();
        Flux<String> result = chatClient.prompt(msg).stream().content()
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("Total time: " + (endTime - startTime) + "ms");
                })
                .doOnError(error -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("Error time: " + (endTime - startTime) + "ms");
                });

        return result;
    }

    //2. 添加进度指示器
    // 由于流式响应的总长度未知，可以基于时间或字符数估算进度：
    @GetMapping("/chat/client/v2.2")
    public Flux<String> hello22(@RequestParam(name="msg",defaultValue = "祖国万岁" )String msg){
        StringBuilder responseBuilder = new StringBuilder();
        long startTime = System.currentTimeMillis();

        Flux<String> result = chatClient.prompt(msg).stream().content()
                .doOnNext(chunk -> {
                    responseBuilder.append(chunk);
                    // 输出当前接收到的字符数
                    System.out.println("Received " + responseBuilder.length() + " characters so far");
                })
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("Completed in " + (endTime - startTime) + "ms");
                    System.out.println("Total characters: " + responseBuilder.length());
                });
        return result;
    }

    //3. 返回带进度信息的包装对象
    // 创建一个包装类来包含内容和进度信息：
    @GetMapping("/chat/client/v2.3")
    public Flux<ProgressResponse> hello23(@RequestParam(name="msg",defaultValue = "祖国万岁" )String msg){
        StringBuilder responseBuilder = new StringBuilder();
        long startTime = System.currentTimeMillis();
        Flux<ProgressResponse> result = chatClient.prompt(msg).stream().content()
                .map(chunk -> {
                    responseBuilder.append(chunk);
                    long currentTime = System.currentTimeMillis();
                    return new ProgressResponse(
                            chunk,
                            null, // 无法准确计算进度，因为不知道总长度
                            currentTime - startTime,
                            false
                    );
                })
                .doOnComplete(() -> {
                    long endTime = System.currentTimeMillis();
                    System.out.println("Completed in " + (endTime - startTime) + "ms");
                });
        return result;
    }

    //4. 使用自定义操作符跟踪进度
    @GetMapping("/chat/client/v2.4")
    public Flux<String> hello24(@RequestParam(name="msg",defaultValue = "祖国万岁" )String msg){
        AtomicLong characterCount = new AtomicLong(0);
        AtomicLong startTime = new AtomicLong(System.currentTimeMillis());
        Flux<String> result = chatClient.prompt(msg).stream().content()
                .doOnNext(chunk -> {
                    long count = characterCount.addAndGet(chunk.length());
                    System.out.println("Chunk #" + count + " received");
                })
                .doOnComplete(() -> {
                    long totalTime = System.currentTimeMillis() - startTime.get();
                    System.out.println("Streaming completed. Total chars: " + characterCount.get() +
                            ", Time: " + totalTime + "ms");
                });
        return result;
    }
}
